#include "reconstruction.hpp"

namespace hms::swe
{

template<typename MeshType>
template<ReconstructionType recType, EdgeRange orien, BlockLocation loc>
void SolverThread<MeshType>::fluxScheme(
	const Block<loc>& block,
	const Array3XsMap& q,
	      Array3XsMap& fluxes,
	const Array1XsMap& h,
	const Ref<const ArrayXXs>& d,
	const Ref<const ArrayXXs>& z,
	const Ref<const ArrayXXs>& v,
	const MeshBase<MeshType>& mesh,
	Index e,
	const Array2i& lr
){
	using RT = ReconstructionType;
	using S = SolverType;
	static constexpr bool
		isUniMesh  { std::is_same_v<MeshType, UniMesh> },
		isRectMesh { std::is_same_v<MeshType, RectMesh> },
		isOrtho    { isUniMesh || isRectMesh };

	// ++(this->nEdges);
	FLUX_TIMER_RESET;

	auto edge { mesh.edge(e) };

	auto isDry { d < S::minWaterDepth };

	/* if both cells are dry, no fluxes occur */
	if ( isDry(0, lr).all() )
		return;

	auto getFieldVals = [&](const auto& field, const Array2i& cells){
		return field(hms::all, cells);
	};

	/* get field values for z and d. also required later for slope source treatment */
	RowArray2s
		zCellLR { getFieldVals(z, lr) },
		dCellLR { getFieldVals(d, lr) };

	EdgeValues ev;

	using namespace reconstruction;

	if ( murillo<S>(dCellLR) && isOrtho && recType != RT::firstOrder ){

		RowArray2s    hCellLR { getFieldVals(h, lr) };
		ArrayNNs<2,2> qCellLR { getFieldVals(q, lr).template bottomRows<2>() };

		if constexpr (recType == RT::secondOrder){
			/* get 2nd second neighbour cells of edge. */
			Array2i llrr {lr};
			if constexpr ( orien == EdgeRange::horizontal ){
				llrr[0] -= block.nxRead();
				llrr[1] += block.nxRead();
			} else {
				llrr[0] -= 1;
				llrr[1] += 1;
			}

			/* h and q are only needed for second order reconstruction */
			RowArray2s
				dCellLLRR { getFieldVals(d, llrr) },
				hCellLLRR { getFieldVals(h, llrr) };

			ArrayNNs<2,2>
				qCellLLRR { getFieldVals(q, llrr).template bottomRows<2>() };

			/* scalar */
			auto [dxQuotL, dxQuotR] = dxQuotient<orien>(mesh, edge.cells(2)[0]);

			RowArray2s psiDH;
			ArrayNNs<2,2> psiQ;

			reconstruction::highOrder<S>(
				ev,
				dCellLR, dCellLLRR, hCellLR, hCellLLRR, qCellLR, qCellLLRR,
				dxQuotL, dxQuotR, solver().settings().limiterFunction
			);
		} else if constexpr (recType == RT::central) {
			reconstruction::central<S>(ev, dCellLR, hCellLR, qCellLR);
		}
	} else {
		reconstruction::firstOrder<S>(
			ev, dCellLR, zCellLR, getFieldVals(v, lr)
		);
	}
	reconstruction::hydrostatic_firstOrder(ev);

	TIME_FLUX(reconstruction);

	if ( (ev.dLR < S::minWaterDepth).all() )
		return;

	/* transformed velocities */
	Array2Ns<2> vRotLR;
	
	if constexpr ( isOrtho ){
		if constexpr ( orien == EdgeRange::horizontal ){
			vRotLR.row(0) =  ev.vLR.row(1);
			vRotLR.row(1) = -ev.vLR.row(0);
		} else {
			vRotLR = ev.vLR;
		}
	} else {
		for (int i{0}; i<2; ++i)
			vRotLR.col(i) = transformCoords(ev.vLR.col(i), edge);
	}

	ArrayNNs<3,2> flux {
		hllc::hllcScheme<S>( ev, vRotLR, edge.normal() )
	};

	TIME_FLUX(hllc);

	/* the slope source term depends on cell centre values and is thus
	 * not equal for both sides of the edge */
	solver().slopeSourceToFlux(
		flux,
		ev.zLR,
		ev.dLR,
		ev.hLR,
		zCellLR,
		dCellLR,
		edge
	);

	TIME_FLUX(slope);

	registerFlux(fluxes, flux.col(0), flux.col(1), edge, lr, mesh);

	TIME_FLUX(registration);
}

template<typename MeshType>
void SolverThread<MeshType>::registerFlux(
	Array3XsMap& fluxes,
	Ref<Array3s> fluxL,
	Ref<Array3s> fluxR,
	const typename MeshType::EdgeType& edge,
	const Array2i& lr,
	const MeshBase<MeshType>& mesh
){
	const auto& lr_mesh { edge.cells(1) };
	fluxL *= solver().time().step() * edge.length() / mesh.cell(lr_mesh[0]).area();
	fluxR *= solver().time().step() * edge.length() / mesh.cell(lr_mesh[1]).area();

	fluxes.col(lr[0]) -= fluxL;
	fluxes.col(lr[1]) += fluxR;
}

} // namespace hms::swe

